#pragma once
#include <boost/asio.hpp>
#include <map>
#include <list>
#include "../../smart_ptr/smart_ptr.hpp"

namespace umtl
{
	//////////////////////////////////////////////////////////////////////////
	//

	typedef boost::asio::ip::tcp tcp;
	typedef boost::asio::ip::udp udp;

	//////////////////////////////////////////////////////////////////////////
	//

	template< typename endpoint_type >
	struct endpoint_key
	{
		endpoint_key() : endpoint_(0) {}

		endpoint_key( endpoint_type & endpoint ) : endpoint_( &endpoint ){}

		endpoint_type const * operator->() const { return endpoint_; }

		endpoint_type * endpoint_;

		bool empty() const { return !endpoint_; }

		friend bool operator<( endpoint_key const & l, endpoint_key const & r )
		{
			if( l.empty() || r.empty() )
				return false;

			std::string ls = l->address().to_string();
			std::string rs = r->address().to_string();

			return ls < rs && ( l->port() < r->port() );
		}
	};

	//////////////////////////////////////////////////////////////////////////
	//

	template< typename protocol >
	struct socket_base
	{
		typedef std::tr1::function< void(boost::system::error_code const&, size_t) >	Handler;
		typedef protocol																protocol;
		typedef typename protocol::socket												socket_type;
		typedef smart_ptr< socket_type >												socket_ptr;
		typedef typename protocol::resolver												resolver_type;
		typedef typename protocol::endpoint												endpoint_type;
		typedef smart_ptr< endpoint_type >												endpoint_ptr;
		typedef std::map< endpoint_key< endpoint_type >, endpoint_ptr >					endpoints_type;

		boost::asio::io_service &	io_service_;
		socket_ptr					socket_;
		boost::asio::strand			strand_;
		resolver_type				resolver_;
		endpoints_type				endpoints_;

		socket_base( boost::asio::io_service & io_service )
			: io_service_(io_service), socket_(new socket_type(io_service)), strand_(io_service), resolver_(io_service) {}

		virtual ~socket_base() = 0 {};

		virtual void listen( unsigned short int /*port*/ ) {}
		virtual void connect( std::string const & host, std::string const & port ) = 0;
		virtual void read( boost::asio::mutable_buffers_1 & buf, Handler handler ) = 0;
		virtual void send( boost::asio::const_buffers_1 & buf, Handler handler ) = 0;
	};

	//////////////////////////////////////////////////////////////////////////
	//

	struct tcp_socket : public socket_base< tcp >
	{
		virtual ~tcp_socket() {}
		tcp_socket( boost::asio::io_service & io_service )
			: socket_base( io_service )
		{}

		virtual void connect( std::string const & host, std::string const & port )
		{ 
			resolver_type::query query(protocol::v4(), host, port);

			resolver_type::iterator connected_endpointer;
			resolver_type::iterator end;

			boost::system::error_code error = boost::asio::error::host_not_found;

			for( resolver_type::iterator endpoint_iter = resolver_.resolve(query)
				; error && endpoint_iter != end; ++endpoint_iter )
			{
				socket_->close();				
				socket_->connect( *endpoint_iter, error );
				connected_endpointer = endpoint_iter;
			}

			if( error )
			{
				throw boost::system::system_error(error);
			}
			else if( connected_endpointer != end )
			{
				endpoint_type * endpoint = new endpoint_type( *connected_endpointer );
				endpoints_.insert( std::make_pair( endpoint_key<endpoint_type>( *endpoint ), endpoint_ptr( endpoint ) ) );
			}
		}

		virtual void read( boost::asio::mutable_buffers_1 & buf, Handler handler ) { socket_->async_read_some( buf, strand_.wrap( handler ) ); }
		virtual void send( boost::asio::const_buffers_1 & buf, Handler handler ) { socket_->async_write_some( buf, strand_.wrap( handler ) ); }
	};

	//////////////////////////////////////////////////////////////////////////
	//

	struct udp_socket : public socket_base< udp >
	{
		typedef std::list< endpoint_type* > reserved_endpoints;
		reserved_endpoints reserved_endpoints_;

		virtual ~udp_socket() {
			std::for_each( reserved_endpoints_.begin(), reserved_endpoints_.end(),
				[ this ]( reserved_endpoints::value_type & endpoint )
				{
					delete endpoint;
				}
			);
		}

		udp_socket( boost::asio::io_service & io_service )
			: socket_base(io_service)
		{}

		virtual void listen( unsigned short int port )
		{
			socket_type * new_socket = new socket_type(io_service_, protocol::endpoint(protocol::v4(), port));

			socket_->close();
			socket_ = new_socket;
		}

		virtual void connect( std::string const & host, std::string const & port )
		{ 
			resolver_type::query query(protocol::v4(), host, port);

			protocol::endpoint receiver_endpoint;

			receiver_endpoint = *resolver_.resolve(query);

			socket_->open( protocol::v4() );

			std::tr1::array<char, 4> send_buf  = { 0 };
			socket_->send_to(boost::asio::buffer(send_buf.data(), send_buf.size()), receiver_endpoint );
		}

		virtual void read( boost::asio::mutable_buffers_1 & buf, Handler handler ) 
		{
			endpoint_type * endpoint = new endpoint_type;

			endpoint->port( 0 );

			socket_->async_receive_from( buf, *endpoint, strand_.wrap( handler ) );

			//endpoints_[*endpoint].reset( endpoint );
			reserved_endpoints_.push_back( endpoint );
		}

		virtual void send( boost::asio::const_buffers_1 & buf, Handler handler ) 
		{
			for( reserved_endpoints::iterator iter = reserved_endpoints_.begin();
				iter != reserved_endpoints_.end(); )
			{
				if( (*iter)->port() )
				{
					endpoints_[ **iter ] = *iter;

					reserved_endpoints_.erase( iter++ );
				}
				else
				{
					++iter;
				}
			}

			std::for_each( endpoints_.begin(), endpoints_.end(), 
				[this, &buf, &handler]( endpoints_type::value_type const & endpoint )
				{
					if( endpoint.second->port() )
						socket_->async_send_to( buf, *endpoint.second, strand_.wrap( handler ) );
				}
			);			
		}
	};

	//////////////////////////////////////////////////////////////////////////
	//

	template< typename protocol >
	struct socket_type_selector {
		typedef socket_base< protocol > type;
	};

	template<>
	struct socket_type_selector< tcp > {
		typedef tcp_socket type;
	};

	template<>
	struct socket_type_selector< udp > {
		typedef udp_socket type;
	};

	//////////////////////////////////////////////////////////////////////////
}